from osgeo import gdal
import numpy as np
import os 
from glob import glob
from math import ceil
import time

# tif 图像的边界
def GetExtent(infile):
    ds = gdal.Open(infile)
    geotrans = ds.GetGeoTransform()
    xsize = ds.RasterXSize
    ysize = ds.RasterYSize
    min_x,max_y = geotrans[0],geotrans[3]
    max_x,min_y = geotrans[0]+xsize*geotrans[1],geotrans[3]+ysize*geotrans[5]
    ds = None
    return min_x,max_y,max_x,min_y

def compress(path, target_path,method="LZW"): 
    """
    使用gdal进行文件压缩，
    LZW方法属于无损压缩，
    效果非常给力，4G大小的数据压缩后只有三十多M
    """
    dataset = gdal.Open(path)
    driver = gdal.GetDriverByName('GTiff')
    driver.CreateCopy(target_path, dataset, strict=1, options=["TILED=YES", "COMPRESS={0}".format(method)])
    del dataset

def RasterMosaic(file_list,outpath):
    Open = gdal.Open
    min_x,max_y,max_x,min_y=GetExtent(file_list[0])
    for infile in file_list:
        minx,maxy,maxx,miny = GetExtent(infile)
        min_x,min_y = min(min_x,minx),min(min_y,miny)
        max_x,max_y = max(max_x,maxx),max(max_y,maxy)
    
    in_ds = Open(file_list[0])
    in_band=in_ds.GetRasterBand(1)
    geotrans = list(in_ds.GetGeoTransform())
    width,height = geotrans[1],geotrans[5]
    columns = ceil((max_x-min_x)/width)#列数
    rows = ceil((max_y-min_y)/(-height))#行数
    
    outfile = outpath+file_list[0][:2]+'.tif' #结果文件名，可自行修改
    driver=gdal.GetDriverByName('GTiff')
    # out_ds=driver.Create(outfile,columns,rows,1,in_band.DataType)
    out_ds=driver.Create(outfile,columns,rows,1,1)
    out_ds.SetProjection(in_ds.GetProjection())
    geotrans[0]=min_x#更正左上角坐标
    geotrans[3]=max_y
    out_ds.SetGeoTransform(geotrans)
    out_band=out_ds.GetRasterBand(1)
    inv_geotrans=gdal.InvGeoTransform(geotrans)

    for in_fn in file_list:
        in_ds=Open(in_fn)
        in_gt=in_ds.GetGeoTransform()
        offset=gdal.ApplyGeoTransform(inv_geotrans,in_gt[0],in_gt[3])
        x,y=map(int,offset)

        data=in_ds.GetRasterBand(1).ReadAsArray()

        data = (data / 256 * 256).astype(np.uint8)

        out_band.WriteArray(data,x,y)#x，y是开始写入时左上角像元行列号
    del in_ds,out_band,out_ds
    return outfile


if __name__ == "__main__":
    # kw = "20190804"
    # kw = "20190822"
    # kw = "20190823"
    kw = ""
    #-- 匹配串
    pattern = f'*{kw}*{kw}*.tif' 
    #-- 该文件夹下存放了待拼接的栅格
    path = r'E:/datasets/spacenet-dataset/spacenet/SN6_buildings/train/AOI_11_Rotterdam/PS-RGB'
    #-- 拼接结果输出文件
    output_path = f'E:/datasets/client-data/mosaic_spn/ps-rgb_{kw}'
    #-- 压缩结果输出文件
    compress_path = f'E:/datasets/client-data/mosaic_spn/compress_ps-rgb_{kw}.tif'
    os.chdir(path)
    raster_list = sorted(glob(pattern), reverse=True) #读取文件夹下所有tif数据
    print("-------------- {} tifs have been selected --------------".format(raster_list.__len__()))
    result = RasterMosaic(raster_list,outpath = output_path ) #拼接栅格
    compress(result,target_path = compress_path) #压缩栅格